;; LUP Decomposition
;; Copyright (C) 2024  Rocks Mazama
;;
;; This program is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <https://www.gnu.org/licenses/>.

(ns mazama.linalg.LUP)

(defn- switch
  "Switch the `i`th and `j`th elements of `coll`, `i` <= `j`."
  [i j coll]
  (if (== i j)
    coll
    (let [a (take i coll)
          b (drop i coll)
          c (take (- j i) b)
          d (drop (- j i) b)]
      (concat a (vector (first d)) (rest c) (vector (first c)) (rest d)))))

(defn- reduce-indexed
  "Reduces a sequence using indexes.
  `f` should be a function of 3 arguments.
  If `coll` contains no entries, returns `init` and `f` is not called.
  Otherwise, returns the result of applying
  `f` to `init`, 0, and the first element of `coll`, then applying
  `f` to that result, 1, and the second element, etc until `coll` is exhausted."
  [f init coll]
  (loop [init init
         n 0
         coll (seq coll)]
    (if coll
      (recur (f init n (first coll)) (inc n) (next coll))
      init)))

(defn- find-pivot-row-index [A]
  (:k (reduce-indexed #(if (or (nil? (:v %1))
                               (> (abs (first %3)) (abs (:v %1))))
                         {:k %2 :v (first %3)}
                         %1)
                      {}
                      A)))

(defn- identity-row [n i]
  (concat (repeat i 0) [1] (repeat (- n i 1) 0 )))

(defn- identity-rows [n]
  (map #(identity-row n %) (range 0 n)))

(defn- rows
  "Convert a matrix into a sequence of rows, each row being a sequence of elements."
  ([{:keys [ncols data]}]
   (rows ncols data))
  ([ncols data]
   (partition ncols data)))

(defn- rows->matrix
  "Create a matrix from a sequence of rows.
  All rows must have the same number of elements otherwise the result
  is undefined."
  [rows]
  {:ncols (count (first rows)) :data (apply concat rows)})

(defn- prelude [n x coll]
  (concat (repeat n x) coll))

(defn- update-L [L l n]
  (map conj L (prelude n 0 (cons 1 l))))

(defn- reducer [[a11 & as]]
  (fn [{:keys [l A]} [u & us]]
    (let [sf (if (zero? a11) 0 (/ u a11))
          a (map #(- %1 (* sf %2)) us as)]
      {:l (conj l sf) :A (conj A a)})))

(defn- decompose [[first-row & rows]]
  (let [initial-val {:l [] :A []}
        val (reduce (reducer first-row) initial-val rows)]
    (assoc val :u first-row)))

(defn LUP
  "Return a LUP decomposition of the matrix `m`."
  [m]
  (let [ncols (:ncols m)
        nrows (if (<= ncols 0) 0 (/ (count (:data m)) ncols))]
    (loop [j 0
           A (rows m)
           L (repeat nrows [])          ;a sequence of rows, built one column at a time
           U []
           P (identity-rows nrows)]
      (if (or (== j ncols) (== j nrows))
        {:L (rows->matrix L)
         :U (rows->matrix U)
         :P (rows->matrix P)}
        (let [k (find-pivot-row-index A)
              A' (switch 0 k A)
              L' (switch j (+ j k) L)
              P' (switch j (+ j k) P)
              d (decompose A')]          
          (recur (inc j)
                 (:A d)
                 (update-L L' (:l d) j)
                 (conj U (prelude j 0 (:u d)))
                 P'))))))
